# Copyright 2011 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utility classes and functions to merge Excel files"""

import StringIO
import sys
import xlrd
import xlutils.filter


class MultiWorkbookReader(xlutils.filter.BaseReader):
    """Reads multiple workbooks.

    Implements get_workbooks(), but not get_filepaths().
    """
    def __init__(self, file_contents, logfile=sys.stdout):
        """Sets to read files specified in 'filenames'.

        Arguments:
          filenames: list of strings representing input file contents.
          logfile: file-like object to write logs.
        """
        self.__file_contents = file_contents
        self.__logfile = logfile

    def get_workbooks(self):
        for i, f in enumerate(self.__file_contents):
            yield (xlrd.open_workbook(file_contents=f,
                                      formatting_info=True,
                                      logfile=self.__logfile),
                   str(i))


class MergedWorkbookWriter(xlutils.filter.BaseWriter):
    """Merges passed workbooks to one workbook."""
    def __init__(self):
        self.__wtbook_name = None

    # Makes the output workbook sticky to the first instance.
    def __setattr__(self, name, value):
        if name == 'wtbook' and self.wtbook is not None:
            pass
        else:
            self.__dict__[name] = value

    # Disallows closing the output workbook implicitly.
    def close(self):
        pass

    def get_workbook(self):
        """Gets the output workbook.
        
        Returns:
          xlwt.Workbook object representing the output.
        """
        return self.wtbook

    def workbook(self, rdbook, wtbook_name):
        self.__wtbook_name = wtbook_name.decode('utf8')
        xlutils.filter.BaseWriter.workbook(self, rdbook, wtbook_name)

    def sheet(self, rdsheet, wtsheet_name):
        xlutils.filter.BaseWriter.sheet(
            self, rdsheet, self.__wtbook_name + '_' + wtsheet_name)


class FirstSheetFilter(xlutils.filter.BaseFilter):
    """Passes only the first sheet to the next processor."""
    def __init__(self):
        self.__first_rdsheet = None
        self.__is_first_sheet = None

    def workbook(self, rdbook, wtbook_name):
        if rdbook.nsheets > 0:
            self.__first_rdsheet = rdbook.sheet_by_index(0)
        self.next.workbook(rdbook, wtbook_name)

    def sheet(self, rdsheet, wtsheet_name):
        self.__is_first_sheet = (rdsheet == self.__first_rdsheet)
        if self.__is_first_sheet:
            self.next.sheet(rdsheet, wtsheet_name)

    def set_rdsheet(self, rdsheet):
        self.__is_first_sheet = (rdsheet == self.__first_rdsheet)
        if self.__is_first_sheet:
            self.next.set_rdsheet(rdsheet)

    def row(self, rdrowx, wtrowx):
        if self.__is_first_sheet:
            self.next.row(rdrowx, wtrowx)

    def cell(self, rdrowx, rdcolx, wtrowx, wtcolx):
        if self.__is_first_sheet:
            self.next.cell(rdrowx, rdcolx, wtrowx, wtcolx)


class MergeBooksFilter(xlutils.filter.BaseFilter):
    """Merges passed books to one book by concatenating sheets."""
    # The current implemation is tricky. Its internal is a sequence of
    # MergedWrokbookWriter, which merges input books, and XLRDReader, which
    # drives the downstream filter with the merged book. This complication is
    # a result of recycling cell style initialization codes in Writer, which
    # is needed for each new book.
    def __init__(self, merged_book_name):
        self.__merged_book_name = merged_book_name
        self.__writer = MergedWorkbookWriter()

    def start(self):
        self.__writer.start()

    def workbook(self, rdbook, wtbook_name):
        self.__writer.workbook(rdbook, wtbook_name)

    def sheet(self, rdsheet, wtsheet_name):
        self.__writer.sheet(rdsheet, wtsheet_name)

    def set_rdsheet(self, rdsheet):
        self.__writer.set_rdsheet(rdsheet)

    def row(self, rdrowx, wtrowx):
        self.__writer.row(rdrowx, wtrowx)
    
    def cell(self, rdrowx, rdcolx, wtrowx, wtcolx):
        self.__writer.cell(rdrowx, rdcolx, wtrowx, wtcolx)

    def finish(self):
        self.__writer.finish()
        buf = StringIO.StringIO()
        self.__writer.get_workbook().save(buf)
        reader = xlutils.filter.XLRDReader(
            xlrd.open_workbook(file_contents=buf.getvalue(),
                               formatting_info=True,
                               logfile=None),
            self.__merged_book_name)
        reader(self.next)


class MergeSheetsFilter(xlutils.filter.BaseFilter):
    """Merges passed sheets to one sheet."""
    def __init__(self, merged_sheet_name, nrows_header=0):
        """Initialization.

        Header rows of the second and later sheets will be skipped.

        Arguments:
          merged_sheet_name: String representing the merged sheet name.
          nrows_header: Integer representing the number of header rows.
        """
        self.__num_processed_sheets = 0
        self.__merged_sheet_name = merged_sheet_name
        self.__nrows_header = nrows_header

        self.__first_rdsheet = None
        self.__nrows_to_skip = None
        self.__row_offsets = {} # Row index offset for each rdsheet
        self.__next_row_offset = 0
        self.__current_row_offset = None

    def __switch_sheet(self, rdsheet):
        """Switches sheet context to 'rdsheet'.

        Arguments:
          rdsheet: xlrd.sheet.Sheet object.
        """
        if self.__first_rdsheet is None:
            self.__first_rdsheet = rdsheet

        # Skips header rows when the sheet is not the first one.
        if rdsheet != self.__first_rdsheet:
            self.__nrows_to_skip = self.__nrows_header
        else:
            self.__nrows_to_skip = 0

        # Registers the sheet row offset if necessary.
        if rdsheet not in self.__row_offsets:
           self.__row_offsets[rdsheet] = self.__next_row_offset
           self.__next_row_offset += (rdsheet.nrows - self.__nrows_header)

        self.__current_row_offset = self.__row_offsets[rdsheet]

    def sheet(self, rdsheet, wtsheet_name):
        self.__switch_sheet(rdsheet)
        if rdsheet == self.__first_rdsheet:
            self.next.sheet(rdsheet, self.__merged_sheet_name)
        else:
            self.next.set_rdsheet(rdsheet)

    def set_rdsheet(self, rdsheet):
        self.__switch_sheet(rdsheet)
        self.next.set_rdsheet(rdsheet)

    def row(self, rdrowx, wtrowx):
        if rdrowx >= self.__nrows_to_skip:
            self.next.row(rdrowx, wtrowx + self.__current_row_offset)

    def cell(self, rdrowx, rdcolx, wtrowx, wtcolx):
        if rdrowx >= self.__nrows_to_skip:
            self.next.cell(rdrowx, rdcolx,
                           wtrowx + self.__current_row_offset, wtcolx)


def MergeFirstSheetsToOneSheet(excel_contents,
                               nrows_header,
                               merged_sheet_name,
                               logfile=sys.stdout):
    """Merges first sheets of Excel files into single-sheet Excel file.

    Arguments:
      excel_contents: strings representing Excel file contents.
      nrows_header: integer representing the number of header rows.
      merged_sheet_name: string representing the merged sheet name.
      logfile: file-like object to write logs.

    Returns:
      string representing the merged Excel file content.
    """
    reader = MultiWorkbookReader(excel_contents, logfile=logfile)
    writer = xlutils.filter.XLWTWriter()
    xlutils.filter.process(reader,
                           FirstSheetFilter(),
                           MergeBooksFilter('merged_book'),
                           MergeSheetsFilter(merged_sheet_name,
                                             nrows_header=nrows_header),
                           writer)
    buf = StringIO.StringIO()
    writer.output[0][1].save(buf)
    return buf.getvalue()


def MergeAllSheetsToOneSheet(excel_contents,
                             nrows_header,
                             merged_sheet_name,
                             logfile=sys.stdout):
    """Merges all sheets of given Excel files into single-sheet Excel file.

    Arguments:
      excel_contents: strings representing Excel file contents.
      nrows_header: integer representing the number of header rows.
      merged_sheet_name: string representing the merged sheet name.
      logfile: file-like object to write logs.

    Returns:
      string representing the merged Excel file content.
    """
    reader = MultiWorkbookReader(excel_contents, logfile=logfile)
    writer = xlutils.filter.XLWTWriter()
    xlutils.filter.process(reader,
                           MergeBooksFilter('merged_book'),
                           MergeSheetsFilter(merged_sheet_name,
                                             nrows_header=nrows_header),
                           writer)
    buf = StringIO.StringIO()
    writer.output[0][1].save(buf)
    return buf.getvalue()
